# -*- coding: utf-8 -*-
"""
An implementation of the policyValueNet in Tensorflow
Tested in Tensorflow 1.14

@author: Liangyue Jia
"""
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable
print('可用gpu数量:',torch.cuda.device_count())  # 可用gpu数量
print('是否可用gpu', torch.cuda.is_available())  # 是否可用gpu

def set_learning_rate(optimizer, lr):
    """Sets the learning rate to the given value"""
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

class Net(nn.Module):
    """policy value network definition"""
    def __init__(self, dim_num, level_num):
        super(Net, self).__init__()

        self.dim_num = dim_num #变量数目
        self.level_num = level_num #变量的水平数
        self.history_steps = 3 #回溯前3步路径

        ## 1. Common Networks Layers
        self.conv1 = nn.Conv2d(self.history_steps, 32, kernel_size=(3, 3), padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=(3, 3), padding=1)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=(3, 3), padding=1)

        ## 2. Actor Networks
        self.act_conv1 = nn.Conv2d(128, 4, kernel_size=(1, 1))
        self.act_fc1 = nn.Linear(4 * dim_num * level_num,
                                 level_num)

        ## 3 Evaluation Networks
        self.val_conv1 = nn.Conv2d(128, 2, kernel_size=(1, 1))
        self.val_fc1 = nn.Linear(2 * dim_num * level_num, 64)
        self.val_fc2 = nn.Linear(64, 1)

    def forward(self, state_input):
        ## 1. Common Networks Layers
        x = F.relu(self.conv1(state_input))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        ## 2. Actor Networks
        x_act = F.relu(self.act_conv1(x))
        x_act_flatten = x_act.view(-1, 4*self.dim_num*self.level_num)
        x_act = F.log_softmax(self.act_fc1(x_act_flatten),dim=0)
        ## 3 Evaluation Networks
        x_val = F.relu(self.val_conv1(x))
        x_val_flatten = x_val.view(-1, 2*self.dim_num*self.level_num)
        x_val = F.relu(self.val_fc1(x_val_flatten))
        x_val = self.val_fc2(x_val)
        return x_act, x_val


class PolicyValueNet():
    def __init__(self, dim_num, level_num, model_file=None):
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.dim_num = dim_num #变量数目
        self.level_num = level_num #变量的水平数
        self.history_steps = 3 #回溯前3步路径
        self.l2_const = 1e-4

        ## 1.构建RL网络结构
        self.policy_value_net = Net(dim_num, level_num)
        self.policy_value_net.to(self.device) #移动模型到gpu或cpu
        ## 2.定义奖励函数

        self.optimizer = optim.Adam(self.policy_value_net.parameters(),
                                    weight_decay=self.l2_const)

        ## 3.存在历史模型则读取
        if model_file:
            net_params = torch.load(model_file)
            self.policy_value_net.load_state_dict(net_params)



    def policy_value(self, state_batch):
        """
        :param state_batch: a batch of states ,输入为一个batch的state数据
        :return: a batch of action probabilities and state values
        """
        state_batch = Variable(torch.FloatTensor(state_batch).to(self.device))
        log_act_probs, value = self.policy_value_net(state_batch)
        # act_probs = np.exp(log_act_probs.data.cpu().numpy()) # 结果数据由cpu处理
        act_probs = torch.exp(log_act_probs).to(self.device) #置于gpu计算
        return act_probs, value


    def policy_value_fn(self, srq):
        """
        :param srq: 问题对象， 即单个局面
        :return: zip:当前状态下可选的（act,prob）列表
                 value: 当前状态下的价值
        """
        depth, legal_levels = srq.get_avaliable_move() #legal_value：[levels , ];
        # legal_levels = torch.FloatTensor(legal_levels).to(self.device)
        current_state = torch.FloatTensor(srq.current_state()).view(
            -1, self.history_steps, self.dim_num, self.level_num).to(self.device)
        log_act_probs, value = self.policy_value_net(current_state) #log_act_probs：[1 X levels]; value:[1 X 1]

        # TODO 现在使用zip，你是不是要改成元祖，链接的是MCTS的expand
        act_probs_zip = zip(legal_levels, log_act_probs[0])  #现在是zip处理
        # act_probs_zip = (legal_levels, log_act_probs[0])  #
        return act_probs_zip, value.item()

    def train_step(self, state_batch, mcts_probs, value_batch, lr):
        """训练过程"""
        # 1:所有数据迁移到gpu或是cpu中
        state_batch = Variable(torch.FloatTensor(state_batch).to(self.device))
        value_batch = Variable(torch.FloatTensor(value_batch).to(self.device))
        mcts_probs = Variable(torch.FloatTensor(mcts_probs.float()).to(self.device))

        # 初始化梯度=0 和学习率
        self.optimizer.zero_grad()
        set_learning_rate(self.optimizer, lr)

        ## forward
        log_act_probs, value = self.policy_value_net(state_batch)
        # 误差表达式：loss = (z - v)^2 - pi^T * log(p) + c||theta||^2
        value_loss = F.mse_loss(value.view(-1), value_batch.view(-1))
        policy_loss = -torch.mean(torch.sum(mcts_probs*log_act_probs, 1))
        loss = value_loss + policy_loss

        ## backward and optimize
        loss.backward()
        self.optimizer.step()

        ##j 计算 policy entropy 用于查看训练效果
        entropy = -torch.mean(
            torch.sum(torch.exp(log_act_probs)*log_act_probs, 1)
        )
        return loss, entropy

    def get_policy_params(self):
        net_params = self.policy_value_net.state_dict()
        return net_params

    def save_model(self, model_file):
        net_params = self.get_policy_params()
        torch.save(net_params, model_file)


if __name__ == '__main__':
    pv_net = PolicyValueNet(8,8)